package org.chorus;

import java.awt.Point;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Map;
import java.util.TreeMap;

import org.gicentre.utils.stat.AbstractChart;
import org.gicentre.utils.stat.XYChart;

import jwo.landserf.structure.RasterMap;
import jwo.landserf.structure.SpatialModel;

import processing.core.PApplet;
import processing.core.PConstants;
import processing.core.PGraphics;
import processing.core.PVector;

//  **********************************************************************************
/** Module for creating and viewing empirical semi-variograms generated from a DEM.
 *  @author Jo Wood
 *  @version 1.3.2, 22nd March, 2013.
 */
//  **********************************************************************************

/* This file is part of the Chorus river bed analysis software package. Chorus is free software:
 * you can redistribute it and/or modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation, either version 3 of the License, or (at your
 * option) any later version.
 * 
 * Chorus is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
 * without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  
 * See the GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License along with this
 * source code (see COPYING included with this source code). If not, see 
 * http://www.gnu.org/licenses/.
 */

public class Variogram implements Modular,Comparable<Modular>
{
	// ----------------------------- Object variables ------------------------------

	private PApplet parent;					// Parent sketch in which this class is embedded.
	private PGraphics graphics;				// Graphics context for output.
	private ArrayList<XYChart> charts;		// Graphical representation of the semivariograms.
	private ArrayList<String> titles;		// Titles of DEMs from which to calculate the semivariograms.
	private boolean showGuides;				// Determines if chart guides are displayed.
	private int order;						// Order number for sorted modules.
	private String units;					// Planimetric units to be labelled on the variogram x-axis.
	private Direction direction, downstreamDir;
	private int colourOffset;
	private float maxLag;					// Maximum permitted lag.
		
	// ------------------------------- Constructors ---------------------------------

	/** Creates a semi-variogram. This version will calculate lags in all directions up to a maximum lag
	 *  of the smaller side of the surface being measured.
	 */
	public Variogram()
	{
		this(Direction.ALL, Direction.LEFT_TO_RIGHT,0);
	}
	
	/** Creates a semi-variogram. This version will calculate lags in all directions up to a maximum lag
	 *  of the smaller side of the surface being measured.
	 *  @param maxLag Maximum lag in grid cell counts or 0 if half the short side of the surface.
	 */
	public Variogram(float maxLag)
	{
		this(Direction.ALL, Direction.LEFT_TO_RIGHT,maxLag);
	}
	
	/** Creates a semi-variogram using lags in the given direction. Maximum permitted lag will be half the
	 *  shorter side of the surface being measured.
	 *  @param direction Direction along which to calculate lags. This can either be DOWNSTREAM or CROSSSTREAM.
	 *  @param downstreamDir Downstream direction in raster. This can be either LEFT_TO_RIGHT or TOP_TO_BOTTOM.
	 */
	public Variogram(Direction direction, Direction downstreamDir)
	{
		this(direction, downstreamDir,0);
	}
	
	/** Creates a semi-variogram using lags in the given direction up to a maximum given lag.
	 *  @param direction Direction along which to calculate lags. This can either be DOWNSTREAM or CROSSSTREAM.
	 *  @param downstreamDir Downstream direction in raster. This can be either LEFT_TO_RIGHT or TOP_TO_BOTTOM.
	 *  @param maxLag Maximum lag in grid cell counts or 0 if half the short side of the surface.
	 */
	public Variogram(Direction direction, Direction downstreamDir, float maxLag)
	{
		this.direction = direction;
		this.downstreamDir = downstreamDir;
		this.maxLag = maxLag;
	}
	
	
	// ---------------------------------- Methods ----------------------------------

	/** Initialises the semi-variogram module.
	 */
	public void setup(PApplet parent)
	{
		this.parent = parent;
		showGuides = false;
		titles = new ArrayList<String>();
		charts = new ArrayList<XYChart>();
		setGraphics(parent.g);
		colourOffset = 0;
		units = "mm";
	}
	
	/** Sets the the graphics context into which all output is directed. This method allows
	 *  output to be redirected to print output, offscreen buffers etc.
	 *  @param graphics New graphics context in which the modular class is embedded.
	 */
	public void setGraphics(PGraphics graphics)
	{
		this.graphics = graphics;
		
		for (AbstractChart chart : charts)
		{
			chart.setGraphics(graphics);
		}
	}
	
	/** Sets the order number for this module. Allows a list of modules to be sorted in some arbitrary order.
	 *  @param order Order number where the smaller the number, the earlier in a sorted list the module will appear.
	 */
	public void setOrder(int order)
	{
		this.order = order;
	}
	
	/** Reports the order number for this module. Allows a list of modules to be sorted in some arbitrary order.
	 *  @return Order number where the smaller the number, the earlier in a sorted list the module will appear.
	 */
	public int getOrder()
	{
		return order;
	}
	
	/** Sets the units labels to be displayed on the horizontal axis of the variogram.
	 *  @param units Units labels (e.g. 'mm' or 'cm').
	 */
	public void setUnits(String units)
	{
		this.units = units;
		
		for (XYChart chart : charts)
		{
			switch (direction)
			{
				case CROSSSTREAM:
					chart.setXAxisLabel("Cross-stream Lag ("+units+")");
					break;
				case DOWNSTREAM:
					chart.setXAxisLabel("Downstream Lag ("+units+")");
					break;
				default:
					chart.setXAxisLabel("Lag (all directions, "+units+")");
					break;
			}
		}
	}

	/** Draws the semi-variogram.
	 *  @param xOrigin Left-hand pixel coordinate of the area in which to draw graph.
	 *  @param yOrigin Top pixel coordinate of the area in which to draw graph.
	 *  @param width Width in pixels of the area in which to draw graph.
	 *  @param height Height in pixels of the area in which to draw graph.
	 */
	@SuppressWarnings("null")
	public void draw(float xOrigin, float yOrigin, float width, float height)
	{
		// Check to see that we have at least one DEM to display		
		if (charts.size() == 0)
		{
			graphics.fill(graphics.color(80));
			graphics.textFont(DisplayUtils.mediumFont);
			graphics.textAlign(PConstants.CENTER, PConstants.CENTER);
			graphics.text("No DEMs available for variogram display",xOrigin+width/2,yOrigin+height/2);
			return;
		}

		PVector blPos=null, trPos = null;
		
		// Maintain a scale with respect to enclosing display area (stops text and lines thickness being to large).
		float graphScale = width/400;
		
		graphics.pushMatrix();
		graphics.translate(xOrigin,yOrigin);
		graphics.scale(graphScale,graphScale);
		graphics.translate(-xOrigin, -yOrigin);

		for (int i=0; i<charts.size(); i++)
		{
			// Draw semivariogram.
			XYChart chart = charts.get(i);
			
			int chartColour;
			if (charts.size() == 1)
			{
				chartColour = parent.color(0);
			}
			else
			{
				chartColour = DisplayUtils.getCategoryPalette().findColour(colourOffset+i);
			}
			
			chart.setLineColour(parent.color(chartColour,100));
			chart.setPointColour(chartColour);
			
			graphics.textFont(DisplayUtils.smallFont);
			chart.draw(xOrigin, yOrigin, width/graphScale,height/graphScale);

			blPos = charts.get(0).getDataToScreen(new PVector(charts.get(0).getMinX(),0));
			trPos = charts.get(0).getDataToScreen(new PVector(charts.get(0).getMaxX(),charts.get(0).getMaxY()));

			// Show title of raster for which semivariogram is displayed.
			graphics.fill(chartColour);
			graphics.textFont(DisplayUtils.smallFont);
graphics.textSize(6);
			graphics.textAlign(PConstants.RIGHT,PConstants.BOTTOM);
			float textHeight = graphics.textAscent()+graphics.textDescent()+2;
			graphics.text(titles.get(i),trPos.x,blPos.y-(charts.size()-i)*textHeight);
		}

		// Show guides if requested
		if (showGuides)
		{
			PVector mouseData = charts.get(0).getScreenToData(new PVector(((parent.mouseX-xOrigin)*graphScale)+xOrigin,
					                                                      ((parent.mouseY-yOrigin)*graphScale)+yOrigin));

			if (mouseData != null)
			{
				graphics.stroke(0,120);
				graphics.strokeWeight(0.5f);
				graphics.line(parent.mouseX,blPos.y,parent.mouseX,trPos.y);
				graphics.line(blPos.x,parent.mouseY,trPos.x,parent.mouseY);
				graphics.textFont(DisplayUtils.smallFont);
				graphics.fill(0,120);
				graphics.textAlign(PConstants.LEFT,PConstants.BOTTOM);
				graphics.text(mouseData.x,parent.mouseX+5,blPos.y);
				graphics.textAlign(PConstants.RIGHT,PConstants.TOP);
				graphics.text(mouseData.y,trPos.x,parent.mouseY+5);
			}
		}
		
		graphics.popMatrix();
	}
	
	
	/** Adds the given DEM to those whose semi-variograms are displayed.
	 *  @param dem DEM to add.
	 */
	public void addDEM(RasterMap dem) 
	{
		if (dem != null)
		{
			buildVariogram1D(dem);
			titles.add(dem.getHeader().getTitle());
		}
	}
	
	/** Sets the colour palette start colour for display.
	 *  @param offset Position in default colour palette in which the first variogram is displayed.
	 */
	public void setColourOffset(int offset)
	{
		this.colourOffset = offset;
	}
	
	/** Sets the maximum semi-covariance value to be shown on the variogram.
	 *  @param maxVariance Maximum semi-covariance.
	 */
	public void setMaxVariance(float maxVariance)
	{
		for (XYChart chart : charts)
		{
			chart.setMaxY(maxVariance);
			chart.transposeAxes(false);
		}
	}
	
	/** Reports the maximum semi-covariance value of the variogram.
	 *  @return Maximum semi-covariance shown by the variogram
	 */
	public float getMaxVariance()
	{
		float maxVar = 0;
		
		for (XYChart chart : charts)
		{
			maxVar = Math.max(maxVar,chart.getMaxY());
		}
		return maxVar;
	}
	
	/** Writes the the lag-semi-variance values out to the given file.
	 *  @param fileName Name of output file.
	 *  @return True if output is written successfully.
	 */
	public boolean writeOutput(String fileName)
	{
		try
		{
			PrintWriter output = new PrintWriter(fileName);
						
			for (int i=0; i<charts.size(); i++)
			{
				// Get semivariogram data.
				XYChart chart = charts.get(i);
				
				float[] lags = chart.getXData();
				float[] semivars = chart.getYData();
				
				output.println("# Lags and standardised semivariance for "+titles.get(i));
				switch (direction)
				{
					case DOWNSTREAM:
						output.println("# Lags calculated in the downstream direction only");
						break;

					case CROSSSTREAM:
						output.println("# Lags calculated in the cross-stream direction only");
						break;
						
					default:
						output.println("# Lags calculated in all directions");
						break;
				}
				
				for (int j=0; j<lags.length; j++)
				{
					output.println(lags[j]+","+semivars[j]);
				}
			}
			
			output.flush();
			output.close();
			System.out.println("Semi-variogram output written to "+fileName);
		}
		catch (IOException e)
		{
			System.out.println("Problem writing Semi-variogram output to "+fileName+"': "+e);
			return false;
		}
		return true;
	}

	// ------------------------ Mouse and keyboard handling ------------------------
	
	/** Activates the guide display when the mouse is pressed.
	 */
	public void mousePressed()
	{
		showGuides = true;
	}

	/** Deactivates the guide display when the mouse is released.
	 */
	public void mouseReleased()
	{
		showGuides = false;
	}

	/** Updates the guide display as the mouse is dragged.
	 */
	public void mouseDragged()
	{
		// Do nothing in this case.
	}
	
	/** Allows this module to be compared with another to determine sort order.
	 *  @param otherModule Other module to compare with this one.
	 */
	public int compareTo(Modular otherModule)
	{
		int comparison = order-otherModule.getOrder();
		if (comparison == 0)
		{
			return hashCode()-otherModule.hashCode();
		}
		return comparison;
	}

	// ------------------------------ Private methods ------------------------------

	/** Creates a 1D semi-variogram from a random sample of raster cell pairs.
	 *  @param dem DEM from which to calculate semi-variogram.
	 */
	private void buildVariogram1D(RasterMap dem)
	{		
		float maxGlobalLag = 0;
		float maxGlobalCovar = 0;
		float maxLagCell = Math.min(dem.getNumRows()/2,dem.getNumCols()/2);
		float gridRes = (dem.getXRes()+dem.getYRes())/2f;

		// Limit the maximum lag by parameter if specified or by the maximum permitted by dimensions of the raster
		if (maxLag > 0) 
		{
			if (direction == Direction.ALL)
			{
				maxLagCell = Math.min(maxLag/gridRes, Math.min(dem.getNumRows()/2,dem.getNumCols()/2));
			}
			else if (direction == Direction.DOWNSTREAM)
			{
				if (downstreamDir==Direction.LEFT_TO_RIGHT)
				{
					maxLagCell = Math.min(maxLag/gridRes,dem.getNumCols()/2);
				}
				else
				{
					maxLagCell = Math.min(maxLag/gridRes,dem.getNumRows()/2);
				}
			}
			else if (direction == Direction.CROSSSTREAM)
			{
				if (downstreamDir==Direction.TOP_TO_BOTTOM)
				{
					maxLagCell = Math.min(maxLag/gridRes,dem.getNumCols()/2);
				}
				else
				{
					maxLagCell = Math.min(maxLag/gridRes,dem.getNumRows()/2);
				}
			}
		}
		
		// For computational speed, use raster cells as lag units but scale to geographic units for graph.
		float minLagCell = 1;
		Map<Float, Float> data = new TreeMap<Float, Float>(); 
		float variance = RasterSummary.calcVariance(dem);
		
		for (float lag=minLagCell; lag<maxLagCell; lag*=1.1)
		{
			Covar covar = new Covar();

			// Find a set of random sample pairs at each lag.
			for (int sample=0; sample<10000; sample++)
			{
				float z1, z2=0;
				Point p1,p2=null;
				boolean foundPair = false;

				do
				{
					// Find first non-null random point
					do
					{
						p1 = new Point((int)(Math.random()*dem.getNumCols()),
									   (int)(Math.random()*dem.getNumRows()));
						z1 = dem.getAttribute(p1.y, p1.x);
					}
					while (z1 == SpatialModel.NO_VALUE);
						
					// Find second non-null point in a random direction and given lag from the first.
					for (int j=0; j<8; j++)		// Allow up to 8 attempts to find a valid neighbour.
					{
						double theta;
						switch (direction)
						{
							case CROSSSTREAM:
								if (downstreamDir == Direction.LEFT_TO_RIGHT)
								{
									theta = Math.random()<0.5?Math.PI:0;
								}
								else
								{
									theta = Math.random()<0.5?Math.PI*0.5:Math.PI*1.5;
								}
								break;
							case DOWNSTREAM:
								if (downstreamDir == Direction.LEFT_TO_RIGHT)
								{
									theta = Math.random()<0.5?Math.PI*0.5:Math.PI*1.5;
								}
								else
								{
									theta = Math.random()<0.5?Math.PI:0;
								}
								break;
							default:
								theta = Math.random()*Math.PI*2;
								break;
						}
						
						p2 = new Point((int)Math.round(p1.x + lag*Math.sin(theta)),
								       (int)Math.round(p1.y + lag*Math.cos(theta)));
						z2 = dem.getAttribute(p2.y, p2.x);
						if ((z2 != SpatialModel.NO_VALUE) && (z2 != SpatialModel.OUT_OF_BOUNDS))
						{
							foundPair = true;
							break;
						}
					}
				}
				while (foundPair == false);
				
				// Add the pair to those at this lag.
				covar.add(z1,z2);
			}
			data.put(new Float(lag),new Float(covar.getSemiCovar()/variance));
		}

		// Pass the variance at each lag to the chart.
		float[] xs = new float[data.size()];
		float[] ys = new float[data.size()];
		int n=0;

		for (Float lag : data.keySet())
		{
			xs[n] = lag.floatValue()*gridRes;
			ys[n] = data.get(lag).floatValue();
			n++;
		}

		XYChart newChart = new XYChart(parent);
		newChart.setPointSize(2);
		newChart.setLineWidth(0.3f);
		newChart.setData(xs, ys);
		newChart.setMinX(0);
		newChart.setMaxX(maxLag);
		newChart.setMinY(0);
		newChart.setMaxY(1);
		
		switch (direction)
		{
			case CROSSSTREAM:
				newChart.setXAxisLabel("Cross-stream Lag ("+units+")");
				break;
			case DOWNSTREAM:
				newChart.setXAxisLabel("Down-stream Lag ("+units+")");
				break;
			default:
				newChart.setXAxisLabel("Lag (all directions, ("+units+")");
				break;
		}
		newChart.setYAxisLabel("Standardised semivariance");
		newChart.showXAxis(true);
		newChart.showYAxis(true);
		if (charts.size() > 0)
		{
			newChart.setDecorations(false);
		}		
		charts.add(newChart);

		
		// Ensure all charts have the same scale
		for (XYChart chart : charts)
		{
			maxGlobalLag   = Math.max(maxGlobalLag, chart.getMaxX());
			maxGlobalCovar = Math.max(maxGlobalCovar, chart.getMaxY());
		}
		
		for (XYChart chart : charts)
		{
			chart.setMaxX(maxGlobalLag);
			chart.setMaxY(maxGlobalCovar);
			chart.transposeAxes(false);
		}
	}
	
	/** Class for representing covariance of a set of values.
	 */
	private class Covar
	{
		private double total;
		private int numSamples;

		public Covar()
		{
			total=0;
			numSamples = 0;
		}

		public void add(double z1, double z2)
		{
			total += (z1-z2)*(z1-z2);
			numSamples++;
		}

		public float getSemiCovar()
		{
			if (numSamples == 0)
			{
				return 0;
			}

			return (float)(total/(2*numSamples));
		}

	}
}

